package com.jianjang.llm.controller;

import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

import com.fasterxml.jackson.databind.ObjectMapper;

import reactor.core.publisher.Flux;

/**
 * 聊天控制器
 * 提供调用 Ollama API 的 REST 接口示例
 * @author zhangjian
 * @link https://docs.spring.io/spring-ai/reference/api/chat/ollama-chat.html
 */
@RestController
@RequestMapping("/api/ollama/")
public class OllamaChatController {

    private static final Logger log = LoggerFactory.getLogger(OllamaChatController.class);

    private final OllamaChatModel chatModel;
    private final ObjectMapper objectMapper;

    @Autowired
    public OllamaChatController(OllamaChatModel chatModel, ObjectMapper objectMapper) {
        this.chatModel = chatModel;
        this.objectMapper = objectMapper;
    }

    @PostMapping("/ai/generate")
    public Map<String,String> generate(@RequestBody Map<String, String> request) {
        String message = request.get("message");
        log.info("[GENERATE] 开始生成，消息: {}", message);
        return Map.of("generation", this.chatModel.call(message));
    }

    @GetMapping("/ai/generateStream")
    public Flux<ChatResponse> generateStream(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        log.info("[STREAM] 开始流式生成，消息: {}", message);
        Prompt prompt = new Prompt(new UserMessage(message));
        return this.chatModel.stream(prompt)
                .doOnSubscribe(subscription -> log.info("[STREAM] 流式响应已订阅"))
                .doOnNext(response -> {
                    // 提取流式响应内容
                    String content = "";
                    if (response.getResult() != null && response.getResult().getOutput() != null) {
                        AssistantMessage output = response.getResult().getOutput();
                        // 尝试多种方式提取内容
                        // 方法1: 尝试通过反射获取 text 属性
                        try {
                            java.lang.reflect.Method getTextMethod = output.getClass().getMethod("getText");
                            content = (String) getTextMethod.invoke(output);
                        } catch (NoSuchMethodException e) {
                            // 方法2: 尝试 toString() 并提取内容
                            String outputStr = output.toString();
                            // 如果 toString() 包含内容，尝试提取
                            if (outputStr != null && !outputStr.isEmpty()) {
                                content = outputStr;
                            }
                        } catch (Exception e) {
                            log.debug("[STREAM] 提取内容失败，尝试使用 toString(): {}", e.getMessage());
                            content = output.toString();
                        }
                    }
                    
                    // 打印流式内容
                    if (content != null && !content.isEmpty()) {
                        // 打印每个流式数据块（不换行，模拟流式效果）
                        System.out.print(content);
                        log.debug("[STREAM] 收到流式数据块: {}", content);
                    }
                    
                    // 同时打印完整的响应对象（用于调试）
                    try {
                        log.debug("[STREAM] 完整响应: {}", objectMapper.writeValueAsString(response));
                    } catch (Exception e) {
                        log.debug("[STREAM] 响应对象: {}", response);
                    }
                })
                .doOnComplete(() -> {
                    System.out.println(); // 换行
                    log.info("[STREAM] 流式响应完成");
                })
                .doOnError(error -> {
                    System.out.println(); // 换行
                    log.error("[STREAM] 流式响应出错: {}", error.getMessage(), error);
                });
    }
}

